#!python
#-*- coding:utf-8 -*-

class Trainer(object):
	def __init__(self):
		super(Trainer, self).__init__()
		self.feature = {}
		self.category = {}
	
	def inc_feature(self, feature, cat):
		self.feature.setdefault(feature, {})
		self.feature[feature].setdefault(cat, 0)
		self.feature[feature][cat] += 1

	def inc_category(self, cat):
		self.category.setdefault(cat, 0)
		self.category[cat] += 1

	def feature_count(self, feature, category):
		if feature in self.feature.keys():
			return float(self.feature[feature].get(category, 0))
		return 0.0	

	def category_count(self, category):
		return float(self.category.get(category, 0))
	
	def training(self, features, category):
		for f in features:
			self.inc_feature(f, category)
		self.inc_category(category)

	def probility(self, feature, category):
		category_count = self.category[category]
		if category_count == 0:
			return 0
		return self.feature_count(feature, category) / category_count

	def weighted_prob(self, feature, category, weight=1.0, assume_prob=0.5):
		prob = self.probility(feature, category)
		total = sum([self.feature_count(feature, c) for c in self.category.keys()])
		return (total*prob + weight*assume_prob) / (total + weight)

	def category_prob(self, category):
		return self.category_count(category) / sum(self.category.values())

	def categorys(self):
		return self.category.keys()

class BayesClassify(object):
	def __init__(self, trainer):
		super(BayesClassify, self).__init__()
		self.trainer = trainer
		self.threshold = {}

	def set_threshold(self, feature, threshold):
		self.threshold[feature] = threshold

	def get_threshold(self, feature):
		return self.threshold.get(feature, 1.0)
	
	def prob(self, features, category):
		my_prob = 1.0
		for feature in features:
			my_prob *= self.trainer.weighted_prob(feature, category)
		my_prob = my_prob * (self.trainer.category_prob(category))
		return my_prob

	def classify(self, features, default="Unkown"):
		best_cat = ""
		max = 0
		probs = {}
		for cat in self.trainer.categorys():
			probs[cat] = self.prob(features, cat)
			if max < probs[cat]:
				max = probs[cat]
				best_cat = cat

		times = self.get_threshold(best_cat)
		for k, v in probs.items():
			if k == best_cat : continue
			if max < v*times: return default

		return best_cat
			

def training_test():
	trainer = Trainer()
	trainer.training(["加盟", "网店"], "bad")
	trainer.training(["网店"], "bad")
	trainer.training(["加盟", "小屋"], "bad")
	trainer.training(["不错", "网店", "小屋"], "good")
	trainer.training(["光临", "小屋"], "good")

	test_text = ["加盟", "网店"]
	prober = BayesClassify(trainer)
	prober.set_threshold("bad", 3)
	cat = prober.classify(test_text)
	print cat

if __name__ == "__main__":
	training_test()
